import logging
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from pinecone import Pinecone, ServerlessSpec
from tqdm import tqdm

# 配置日志格式，包含时间、级别和消息
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# 初始化Pinecone客户端
pinecone = Pinecone(api_key="pcsk_4zcKEp_P5uHzGJS6Eq9UhAcqWoiFv6seAxnnxPhFb8yDkHykkHKsCNAiBMvr1TH9Nkzz3f")

# 索引名称
index_name = "mnist-index"

################### 检查、清除之前创建的索引 ###################
# 获取现有索引列表
existing_indexes = pinecone.list_indexes()

# 检查索引是否存在，如果存在就删除
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已成功删除。")
else:
    logging.info(f"索引 '{index_name}' 不存在，将创建新索引。")
################### 检查、清除之前创建的索引 ###################

# 创建新索引
logging.info(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")

# 加载MNIST数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 划分训练集和测试集（80%训练，20%测试）
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 将训练数据转换为Pinecone可接受的格式
vectors = []
for i in range(len(X_train)):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))

# 定义批处理大小
batch_size = 1000

# 使用tqdm显示上传进度条
for i in tqdm(range(0, len(vectors), batch_size), desc="上传数据到Pinecone"):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)

# 记录日志：成功创建索引，并上传了1437条数据
logging.info("成功创建索引，并上传了1437条数据")
print(f"成功创建索引，并上传了{len(vectors)}条数据")

# 测试阶段：使用20%的数据进行查询
correct_predictions = 0
total_queries = len(X_test)

# 使用tqdm显示测试进度条
for i in tqdm(range(total_queries), desc="测试k=11的准确率"):
    query_data = X_test[i].tolist()
    results = index.query(
        vector=query_data,
        top_k=11,
        include_metadata=True
    )
    labels = [match['metadata']['label'] for match in results['matches']]
    final_prediction = max(set(labels), key=labels.count)  # 投票机制
    if final_prediction == y_test[i]:
        correct_predictions += 1

# 计算准确率
accuracy = correct_predictions / total_queries
logging.info(f"当k=11时，准确率为: {accuracy:.4f}")
print(f"当k=11时，准确率为: {accuracy:.4f}")